import numpy as np

from .Position import Position, PositionType


class DistanceMatrix:
    def __init__(self, data):
        self._remote_remote = data["remote_remote"]
        self._remote_boarding = data["remote_boarding"]
        self._arrival_remote = data["arrival_remote"]
        self._arrival_boarding = data["arrival_boarding"]
        self._rest_remote = data["rest_remote"]
        self._rest_boarding = data["rest_boarding"]
        self._rest_arrival = data["rest_arrival"]
    
    def get_distance(self, 
            first_pos: Position, 
            second_pos: Position) -> int:

        # ====================
        # 远机位 -> 远机位
        if first_pos.pos_type == PositionType.Remote and \
                second_pos.pos_type == PositionType.Remote:
            distance = self._remote_remote[second_pos.name][first_pos.name]
            if not np.isnan(distance):
                return distance
            else:
                return self._remote_remote[first_pos.name][second_pos.name]

        # 远机位 -> 登机口  
        if first_pos.pos_type == PositionType.Remote and \
                second_pos.pos_type == PositionType.Boarding:
            return self._remote_boarding[second_pos.name][first_pos.name]

        # 远机位 -> 到达口  
        if first_pos.pos_type == PositionType.Remote and \
                second_pos.pos_type == PositionType.Arrival:
            return self._arrival_remote[first_pos.name][second_pos.name]
        
        # 远机位 -> 休息区  
        if first_pos.pos_type == PositionType.Remote and \
                second_pos.pos_type == PositionType.Rest:
            return self._rest_remote[first_pos.name][second_pos.name]

        # ====================
        # 登机口 -> 远机位
        if first_pos.pos_type == PositionType.Boarding and \
                second_pos.pos_type == PositionType.Remote:
            return self._remote_boarding[first_pos.name][second_pos.name]

        # 登机口 -> 登机口  
        if first_pos.pos_type == PositionType.Boarding and \
                second_pos.pos_type == PositionType.Boarding:
            return 0

        # 登机口 -> 到达口  
        if first_pos.pos_type == PositionType.Boarding and \
                second_pos.pos_type == PositionType.Arrival:
            return self._arrival_boarding[first_pos.name][second_pos.name]
        
        # 登机口 -> 休息区  
        if first_pos.pos_type == PositionType.Boarding and \
                second_pos.pos_type == PositionType.Rest:
            return self._rest_boarding[first_pos.name][second_pos.name]

        # ====================
        # 到达口 -> 远机位
        if first_pos.pos_type == PositionType.Arrival and \
                second_pos.pos_type == PositionType.Remote:
            return self._arrival_remote[second_pos.name][first_pos.name]

        # 到达口 -> 登机口
        if first_pos.pos_type == PositionType.Arrival and \
                second_pos.pos_type == PositionType.Boarding:
            return self._arrival_boarding[second_pos.name][first_pos.name]

        # 到达口 -> 到达口
        if first_pos.pos_type == PositionType.Arrival and \
                second_pos.pos_type == PositionType.Arrival:
            return 0

        # 到达口 -> 休息区
        if first_pos.pos_type == PositionType.Arrival and \
                second_pos.pos_type == PositionType.Rest:
            return self._rest_arrival[first_pos.name][second_pos.name]

        # ====================
        # 休息区 -> 远机位
        if first_pos.pos_type == PositionType.Rest and \
                second_pos.pos_type == PositionType.Remote:
            return self._rest_remote[second_pos.name][first_pos.name]

        # 休息区 -> 登机口
        if first_pos.pos_type == PositionType.Rest and \
                second_pos.pos_type == PositionType.Boarding:
            return self._rest_boarding[second_pos.name][first_pos.name]

        # 休息区 -> 到达口
        if first_pos.pos_type == PositionType.Rest and \
                second_pos.pos_type == PositionType.Arrival:
            return self._rest_arrival[second_pos.name][first_pos.name]

        # 休息区 -> 休息区
        if first_pos.pos_type == PositionType.Rest and \
                second_pos.pos_type == PositionType.Rest:
            return 0